# utils/data_augmenter.py
import numpy as np


class DataAugmenter:
    """数据增强工具 - 用于小数据量情况"""

    @staticmethod
    def augment_sequences(sequences, labels, augmentation_factor=2):
        """增强序列数据"""
        if len(sequences) == 0:
            return sequences, labels

        augmented_sequences = []
        augmented_labels = []

        for sequence, label in zip(sequences, labels):
            # 原始序列
            augmented_sequences.append(sequence)
            augmented_labels.append(label)

            # 数据增强
            for i in range(augmentation_factor):
                # 添加噪声
                noise = np.random.normal(0, 0.01, sequence.shape)
                augmented_seq = sequence + noise
                augmented_sequences.append(augmented_seq)
                augmented_labels.append(label)

                # 时间偏移
                if len(sequence) > 5:
                    shift = np.random.randint(1, 3)
                    augmented_seq = np.roll(sequence, shift, axis=0)
                    augmented_sequences.append(augmented_seq)
                    augmented_labels.append(label)

        return np.array(augmented_sequences), np.array(augmented_labels)

    @staticmethod
    def create_synthetic_data(original_sequences, original_labels, num_synthetic=10):
        """创建合成数据"""
        if len(original_sequences) == 0:
            return original_sequences, original_labels

        synthetic_sequences = []
        synthetic_labels = []

        for sequence, label in zip(original_sequences, original_labels):
            for _ in range(num_synthetic):
                # 基于原始序列创建变体
                synthetic_seq = sequence.copy()

                # 随机缩放
                scale = np.random.uniform(0.8, 1.2)
                synthetic_seq *= scale

                # 随机旋转（简单的轴交换）
                if np.random.random() > 0.5:
                    synthetic_seq = synthetic_seq[:, [1, 0, 2]]  # 交换x和y

                synthetic_sequences.append(synthetic_seq)
                synthetic_labels.append(label)

        return np.array(synthetic_sequences), np.array(synthetic_labels)